package com.czk;

import javassist.CannotCompileException;
import javassist.ClassPool;
import javassist.CtBehavior;
import javassist.CtClass;
import javassist.expr.ExprEditor;
import javassist.expr.MethodCall;
import lombok.SneakyThrows;

import java.lang.instrument.ClassFileTransformer;
import java.lang.instrument.IllegalClassFormatException;
import java.security.ProtectionDomain;
import java.util.HashSet;
import java.util.Set;

/**
 * 是一种探针技术可以通过 premain 方法，在类加载的过程中给指定的方法进行字节码增强。其实你的每一个类最终都是字节码指令的执行，
 * 而这种增强后的方法就可以输出我们想要的信息。这就相当于你硬编码时候输出了一些方法的耗时，日志等信息。
 * @Author:ChenZhangKun
 * @Date: 2021/8/5 11:52
 */
public class MyMonitorTransformer implements ClassFileTransformer {
    // 定义集合
    private static final Set<String> classNameSet = new HashSet<>();

    static {
        classNameSet.add("com.czk.ApiTest");
    }

    @SneakyThrows
    @Override
    public byte[] transform(ClassLoader loader, String className, Class<?> classBeingRedefined, ProtectionDomain protectionDomain, byte[] classfileBuffer) throws IllegalClassFormatException {
        // 拿到当前类
        String currentClass = className.replaceAll("/", ".");
        // 看是否包含
        if (!classNameSet.contains(currentClass)) {
            // 不要增强
            return null;
        }
        // 转换
        System.out.println("transform【" + currentClass + "】");
        CtClass ctClass = ClassPool.getDefault().get(currentClass);
        CtBehavior[] declaredBehaviors = ctClass.getDeclaredBehaviors();
        for (CtBehavior declaredBehavior : declaredBehaviors) {
            // 增强每个方法
            enhanceMethod(declaredBehavior);
        }
        return ctClass.toBytecode();
    }

    private void enhanceMethod(CtBehavior declaredBehavior) throws CannotCompileException {
        if (declaredBehavior.isEmpty()) {
            return;
        }
        // 拿到名称
        String name = declaredBehavior.getName();
        if ("main".equalsIgnoreCase(name)) {
            // 主方法
            return;
        }
        StringBuilder source = new StringBuilder();
        // 前置增强: 打入时间戳
        // 保留原有的代码处理逻辑
        source.append("{")
                // 前置增强: 打入时间戳
                .append("long start = System.nanoTime();\n")
                // 调用原有代码，类似于method();($$)表示所有的参数
                .append("$_ = $proceed($$);\n")
                .append("System.out.print(\"method:[")
                .append(name).append("]\");").append("\n")
                // 后置增强，计算输出方法执行耗时
                .append("System.out.println(\" cost:[\" +(System.nanoTime() - start)+ \"ns]\");")
                .append("}");
        ExprEditor exprEditor = new ExprEditor() {
            @Override
            public void edit(MethodCall m) throws CannotCompileException {
                m.replace(source.toString());
            }
        };
        declaredBehavior.instrument(exprEditor);
    }
}
